package com.example.common;

import com.jcraft.jsch.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;

public class SSHUtil {
    private static final Logger logger = LoggerFactory.getLogger(SSHUtil.class);
    
    // 跳板机配置
    private static final String JUMP_SERVER_HOST = "10.77.18.193";  // 跳板机IP
    private static final String JUMP_SERVER_USER = "root";          // 跳板机用户名
    private static final String JUMP_SERVER_PASSWORD = "sxczt@Boss2023#1";  // 跳板机密码
    private static final String SHARED_DIR = "/tmp/tomcat_files";   // 共享目录

    public static String executeCommand(String targetHost, String command) {
        JSch jsch = new JSch();
        Session session = null;
        ChannelExec channel = null;
        
        try {
            // 如果目标主机与跳板机相同，直接在本机执行
            if (targetHost.equals(JUMP_SERVER_HOST)) {
                session = jsch.getSession(JUMP_SERVER_USER, JUMP_SERVER_HOST, 22);
                session.setPassword(JUMP_SERVER_PASSWORD);
            } else {
                // 通过跳板机执行命令，使用跳板机密码
                session = jsch.getSession(JUMP_SERVER_USER, JUMP_SERVER_HOST, 22);
                session.setPassword(JUMP_SERVER_PASSWORD);
                command = String.format("ssh -o StrictHostKeyChecking=no %s@%s '%s'", 
                    JUMP_SERVER_USER, targetHost, command.replace("'", "'\\''"));
            }
            
            session.setConfig("StrictHostKeyChecking", "no");
            session.connect();

            channel = (ChannelExec) session.openChannel("exec");
            channel.setCommand(command);

            ByteArrayOutputStream responseStream = new ByteArrayOutputStream();
            channel.setOutputStream(responseStream);
            channel.connect();

            while (channel.isConnected()) {
                Thread.sleep(100);
            }

            return new String(responseStream.toByteArray(), StandardCharsets.UTF_8);
        } catch (Exception e) {
            logger.error("Error executing command: {}", e.getMessage());
            throw new RuntimeException("Error executing command: " + e.getMessage());
        } finally {
            if (channel != null) {
                channel.disconnect();
            }
            if (session != null) {
                session.disconnect();
            }
        }
    }

    public static byte[] downloadFile(String targetHost, String filePath) {
        JSch jsch = new JSch();
        Session session = null;
        ChannelSftp channel = null;
        
        try {
            // 连接跳板机
            session = jsch.getSession(JUMP_SERVER_USER, JUMP_SERVER_HOST, 22);
            session.setPassword(JUMP_SERVER_PASSWORD);
            session.setConfig("StrictHostKeyChecking", "no");
            session.connect();

            channel = (ChannelSftp) session.openChannel("sftp");
            channel.connect();

            // 如果目标主机与跳板机不同，通过共享目录传输
            if (!targetHost.equals(JUMP_SERVER_HOST)) {
                String fileName = filePath.substring(filePath.lastIndexOf('/') + 1);
                String sharedFilePath = SHARED_DIR + "/" + fileName;
                
                // 创建共享目录（如果不存在）
                executeCommand(JUMP_SERVER_HOST, "mkdir -p " + SHARED_DIR);
                
                // 使用跳板机密码从目标主机复制文件到共享目录
                String scpCommand = String.format("scp -o StrictHostKeyChecking=no %s@%s:'%s' %s", 
                    JUMP_SERVER_USER, targetHost, filePath, sharedFilePath);
                executeCommand(JUMP_SERVER_HOST, scpCommand);
                
                // 更新文件路径为共享目录中的路径
                filePath = sharedFilePath;
            }

            // 下载文件
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            InputStream inputStream = channel.get(filePath);
            
            byte[] buffer = new byte[1024];
            int len;
            while ((len = inputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, len);
            }

            // 如果使用了共享目录，清理临时文件
            if (!targetHost.equals(JUMP_SERVER_HOST)) {
                executeCommand(JUMP_SERVER_HOST, "rm -f " + filePath);
            }

            return outputStream.toByteArray();
        } catch (Exception e) {
            logger.error("Error downloading file: {}", e.getMessage());
            throw new RuntimeException("Error downloading file: " + e.getMessage());
        } finally {
            if (channel != null) {
                channel.disconnect();
            }
            if (session != null) {
                session.disconnect();
            }
        }
    }

    public static void uploadFile(String targetHost, String fileName, byte[] fileContent) {
        JSch jsch = new JSch();
        Session session = null;
        ChannelSftp channel = null;
        
        try {
            // 连接跳板机
            session = jsch.getSession(JUMP_SERVER_USER, JUMP_SERVER_HOST, 22);
            session.setPassword(JUMP_SERVER_PASSWORD);
            session.setConfig("StrictHostKeyChecking", "no");
            session.connect();

            channel = (ChannelSftp) session.openChannel("sftp");
            channel.connect();

            if (targetHost.equals(JUMP_SERVER_HOST)) {
                // 如果是跳板机本身，直接上传到目标路径
                channel.put(new java.io.ByteArrayInputStream(fileContent), fileName);
            } else {
                // 如果是其他主机，通过共享目录传输
                // 创建共享目录
                executeCommand(JUMP_SERVER_HOST, "mkdir -p " + SHARED_DIR);
                String tempFile = SHARED_DIR + "/" + fileName;

                // 上传文件到共享目录
                channel.put(new java.io.ByteArrayInputStream(fileContent), tempFile);

                // 使用跳板机密码将文件从共享目录复制到目标主机
                String scpCommand = String.format("scp -o StrictHostKeyChecking=no %s %s@%s:/tmp/", 
                    tempFile, JUMP_SERVER_USER, targetHost);
                executeCommand(JUMP_SERVER_HOST, scpCommand);

                // 清理临时文件
                executeCommand(JUMP_SERVER_HOST, "rm -f " + tempFile);
            }
        } catch (Exception e) {
            logger.error("Error uploading file: {}", e.getMessage());
            throw new RuntimeException("Error uploading file: " + e.getMessage());
        } finally {
            if (channel != null) {
                channel.disconnect();
            }
            if (session != null) {
                session.disconnect();
            }
        }
    }
} 